package main

import (
	"bufio"
	"encoding/base64"
	"errors"
	"fmt"
	"gitee.com/kaylee595/gkaylee/kerr"
	proxyPool "gitee.com/kaylee595/proxy-pool"
	"gitee.com/kaylee595/proxy-pool/source"
	"github.com/patrickmn/go-cache"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"
	"tunnelProxy/config"
	"tunnelProxy/proxySourceColonSeparated"
)

var (
	cachePool   *cache.Cache
	proxyIPPool *proxyPool.Pool
)

func init() {
	config.C = kerr.Must(config.Load())
	if config.C.ProxyExpired == 0 {
		config.C.ProxyExpired = 5 * time.Minute
	}
	if config.C.MaxUse == 0 {
		config.C.MaxUse = 1
	}
	cachePool = cache.New(config.C.ProxyExpired, config.C.ProxyExpired+5*time.Minute)
	s := newFetcher(config.C.Url, config.C.Username, config.C.Password, config.C.ProxyExpired)
	proxyIPPool = proxyPool.New(s, config.C.MaxUse)
}

func newFetcher(url_, dUsername, dPassword string, dExpired time.Duration) source.Fetcher {
	//var options []general.Option
	//if dUsername != "" && dPassword != "" {
	//	options = append(options, general.WithAuth(url.UserPassword(dUsername, dPassword)))
	//} else if dUsername != "" {
	//	options = append(options, general.WithAuth(url.User(dUsername)))
	//}
	//return general.New(url_, general.WithExpire(dExpired))

	var options []proxySourceColonSeparated.Option
	if dUsername != "" && dPassword != "" {
		options = append(options, proxySourceColonSeparated.WithAuth(url.UserPassword(dUsername, dPassword)))
	} else if dUsername != "" {
		options = append(options, proxySourceColonSeparated.WithAuth(url.User(dUsername)))
	}
	return proxySourceColonSeparated.New(url_, proxySourceColonSeparated.WithExpire(dExpired))
}

var errEmptyProxyAuthorization = errors.New("empty ProxyAuthorization")
var errNotExistSession = errors.New("not exist session")

func getSession(r *http.Request) (string, error) {
	auth := r.Header.Get("Proxy-Authorization")
	if auth == "" {
		return "", errEmptyProxyAuthorization
	}
	kind, info, ok := strings.Cut(auth, " ")
	if ok {
		switch kind {
		case "Basic":
			decodeInfo, err := base64.StdEncoding.DecodeString(info)
			if err != nil {
				return "", err
			}
			user, _, _ := strings.Cut(string(decodeInfo), ":")
			if user != "" {
				user, err = url.QueryUnescape(user)
				if err != nil {
					return "", err
				}
				fields := strings.FieldsFunc(user, func(r rune) bool { return r == '-' })
				var session string
				for i, field := range fields {
					field = strings.TrimSpace(field)
					if field == "" {
						continue
					}
					if field == "session" {
						if i+1 >= len(fields) {
							break
						}
						session = fields[i+1]
						break
					}
				}
				if session == "" {
					return "", errNotExistSession
				}
				return session, nil
			}
		default:
			return "", fmt.Errorf("invalid auth type: %s", kind)
		}
	}
	return "", errors.New("invalid ProxyAuthorization")
}

func getUpstreamConn(r *http.Request) (net.Conn, *url.Userinfo, error) {
	var upstream *url.URL
	session, err := getSession(r)
	if err != nil {
		if !errors.Is(err, errEmptyProxyAuthorization) && !errors.Is(err, errNotExistSession) {
			return nil, nil, err
		}
		node, err := proxyIPPool.Get(r.Context())
		if err != nil {
			return nil, nil, err
		}
		upstream = node.URL()
		slog.Info("代理获取1", "proxy", upstream.String())
	} else {
		item, ok := cachePool.Get(session)
		if !ok {
			node, err := proxyIPPool.Get(r.Context())
			if err != nil {
				return nil, nil, err
			}
			upstream = node.URL()
			cachePool.SetDefault(session, upstream)
			slog.Info("代理获取2", "proxy", upstream.String(), "session", session)
		} else {
			upstream = item.(*url.URL)
			slog.Info("代理获取3", "proxy", upstream.String(), "session", session)
		}
	}

	switch upstream.Scheme {
	case "http":
		conn, err := net.Dial("tcp", upstream.Host)
		if err != nil {
			return nil, nil, err
		}
		return conn, upstream.User, nil
	default:
		return nil, nil, fmt.Errorf("unsupported scheme: %s", upstream.Scheme)
	}
}

var ErrUpstreamConnectFailed = errors.New("上游CONNECT失败")

func getUpstreamConnAndConnect(r *http.Request) (net.Conn, error) {
	upstreamConn, auth, err := getUpstreamConn(r)
	if err != nil {
		return nil, fmt.Errorf("获取上游连接失败: %w", err)
	}
	_ = upstreamConn.SetDeadline(time.Now().Add(1500 * time.Millisecond))

	// 发送 CONNECT 请求到上游代理
	if auth == nil {
		_, err = fmt.Fprintf(upstreamConn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", r.Host, r.Host)
	} else {
		_, err = fmt.Fprintf(upstreamConn,
			"CONNECT %s HTTP/1.1\r\nHost: %s\r\nProxy-Authorization: Basic %s\r\n\r\n",
			r.Host, r.Host, base64.StdEncoding.EncodeToString([]byte(auth.String())),
		)
	}
	if err != nil {
		_ = upstreamConn.Close()
		return nil, fmt.Errorf("发送CONNECT到上游失败, host: %s, error: %w", r.Host, err)
	}

	// 读取响应，判断是否成功
	resp, err := http.ReadResponse(bufio.NewReader(upstreamConn), r)
	if err != nil {
		_ = upstreamConn.Close()
		return nil, fmt.Errorf("从上游读取响应体失败, host: %s, err: %w", r.Host, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_ = upstreamConn.Close()
		return nil, fmt.Errorf("%w, host: %s, statusCode: %d",
			ErrUpstreamConnectFailed, upstreamConn.RemoteAddr().String(), resp.StatusCode)
	}

	_ = upstreamConn.SetDeadline(time.Time{})
	slog.Info("上游客户端地址", "addr", upstreamConn.RemoteAddr().String())
	return upstreamConn, nil
}

// 处理 CONNECT 方法（用于 HTTPS）
func handleConnect(w http.ResponseWriter, r *http.Request) {
	var upstreamConn net.Conn
	var err error
	for range 8 {
		upstreamConn, err = getUpstreamConnAndConnect(r)
		if err != nil {
			if !errors.Is(err, os.ErrDeadlineExceeded) && !errors.Is(err, ErrUpstreamConnectFailed) {
				slog.Error("获取上游连接失败", "err", err)
				http.Error(w, "获取上游连接失败, err: "+err.Error(), http.StatusServiceUnavailable)
				return
			}
			session, err := getSession(r)
			if err == nil {
				cachePool.Delete(session)
			}
			continue // 超时的重新获取代理连接
		}
		break
	}
	if upstreamConn == nil {
		slog.Error("多轮重试后依旧无法获取上游链接", "err", err)
		http.Error(w, "获取上游连接失败, err: "+err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer upstreamConn.Close()

	// Hijack 客户端连接
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		slog.Error("Hijacking not supported")
		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
		return
	}
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		slog.Error("Hijack failed", "err", err)
		http.Error(w, "Hijack failed: "+err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer clientConn.Close()

	// 告诉客户端隧道建立成功
	_, _ = clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))

	// 双向转发数据（TCP隧道）
	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, err := io.Copy(upstreamConn, clientConn)
		if err != nil {
			slog.Error("从客户端转发数据到上游出错", "err", err)
		}
	}()
	go func() {
		defer wg.Done()
		_, err := io.Copy(clientConn, upstreamConn)
		if err != nil {
			slog.Error("从上游转发数据到客户端出错", "err", err)
		}
	}()
	wg.Wait()
}

// 处理普通 HTTP 请求
func handleHTTP(w http.ResponseWriter, r *http.Request) {
	upstreamConn, auth, err := getUpstreamConn(r)
	if err != nil {
		http.Error(w, "获取上游连接失败: "+err.Error(), http.StatusBadRequest)
		return
	}
	defer upstreamConn.Close()

	// 将客户端请求写入上游连接
	rr, err := http.NewRequest(r.Method, r.URL.String(), r.Body)
	if err != nil {
		http.Error(w, "创建请求失败: "+err.Error(), http.StatusServiceUnavailable)
		return
	}
	rr.Header = r.Header
	if auth != nil {
		rr.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(auth.String())))
	}
	err = rr.WriteProxy(upstreamConn)
	if err != nil {
		http.Error(w, "转发请求到上游失败: "+err.Error(), http.StatusServiceUnavailable)
		return
	}

	// 从上游读取响应并写回客户端
	resp, err := http.ReadResponse(bufio.NewReader(upstreamConn), r)
	if err != nil {
		http.Error(w, "读取上游响应失败: "+err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	// 将上游响应写回客户端
	header := w.Header()
	for key, vv := range resp.Header {
		for _, v := range vv {
			header.Add(key, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

// 代理处理器，根据方法转发到不同处理函数
func handleProxy(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		handleConnect(w, r)
	} else {
		handleHTTP(w, r)
	}
}

// todo: IP使用到本地的问题目前使用SetCommonRetryCount(0), 不知道能不能解决, 需要进一步测试.
// https测试完成之后再测试http, 都没有问题就可以正式使用了.
func main() {
	err := http.ListenAndServe(":9000", http.HandlerFunc(handleProxy))
	if err != nil {
		panic(err)
	}
}
